package viewdata

// TreeData 定义树节点接口
type TreeData[ID any] interface {
	GetId() ID
	GetParentId() ID
	GetChildren() []TreeData[ID]
	SetChildren([]TreeData[ID])
	IsLeaf() bool
}

// BuildTree 构建树形结构
func BuildTree[ID comparable](nodes []TreeData[ID], root TreeData[ID]) []TreeData[ID] {
	if len(nodes) == 0 {
		return []TreeData[ID]{}
	}

	if root.GetChildren() == nil { // 确保父节点的children不为nil
		root.SetChildren([]TreeData[ID]{})
	}

	// 使用map缓存所有节点，key为节点ID，value为节点对象
	nodeMap := make(map[ID]TreeData[ID])
	nodeMap[root.GetId()] = root
	// 第一遍循环：将所有非叶子节点放入map，并找到根节点
	for _, node := range nodes {
		//if !node.IsLeaf() {
		node.SetChildren([]TreeData[ID]{})
		nodeMap[node.GetId()] = node
		//}
	}

	// 第二遍循环：构建父子关系
	for _, node := range nodes {
		parentID := node.GetParentId()

		parent, exists := nodeMap[parentID]
		if exists {
			parent.SetChildren(append(parent.GetChildren(), node))
		}
	}

	return root.GetChildren()
}
